package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// GetUSBDeviceIDs 获取Windows系统上所有连接的USB存储设备的设备ID
func GetUSBDeviceIDs() ([]string, error) {
	// 使用PowerShell命令获取USB存储设备信息
	// DriveType=2 表示可移动磁盘（通常是USB设备）
	psCommand := `Get-WmiObject -Class Win32_LogicalDisk -Filter "DriveType=2" | Select-Object -ExpandProperty DeviceID`
	cmd := exec.Command("powershell", "-Command", psCommand)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("执行PowerShell命令失败: %v, 输出: %s", err, output)
	}

	// 处理命令输出
	result := []string{}
	lines := strings.Split(string(output), "\n")

	// 提取设备ID
	for _, line := range lines {
		deviceID := strings.TrimSpace(line)
		if len(deviceID) > 0 {
			result = append(result, deviceID)
		}
	}

	return result, nil
}

// DeleteAllFilesOnDevice 删除设备根目录下的所有文件和文件夹（保留设备本身）
func DeleteAllFilesOnDevice(deviceID string) error {
	// 确保设备ID格式正确
	if !strings.HasSuffix(deviceID, "\\") {
		deviceID = deviceID + "\\"
	}

	// 获取设备根目录下的所有项目
	entries, err := os.ReadDir(deviceID)
	if err != nil {
		return fmt.Errorf("读取设备根目录失败: %v", err)
	}

	// 逐个删除所有项目
	for _, entry := range entries {
		itemPath := filepath.Join(deviceID, entry.Name())
		if err := os.RemoveAll(itemPath); err != nil {
			return fmt.Errorf("删除项目 %s 失败: %v", itemPath, err)
		}
	}

	return nil
}

// CopyFolderToDevice 将源文件夹复制到目标设备的根目录
func CopyFolderToDevice(sourceFolder, deviceID string) error {
	// 构建目标路径
	targetFolder := filepath.Join(deviceID, "private")

	// 如果目标文件夹已存在，则删除它
	if _, err := os.Stat(targetFolder); err == nil {
		if err := os.RemoveAll(targetFolder); err != nil {
			return fmt.Errorf("删除现有文件夹失败: %v", err)
		}
	}

	// 使用xcopy命令复制文件夹（/E 复制所有子目录，包括空目录；/I 假设目标是目录；/Y 不提示覆盖）
	cmd := exec.Command("cmd.exe", "/C", "xcopy", "/E", "/I", "/Y", sourceFolder, targetFolder)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("复制文件夹失败: %v, 输出: %s", err, output)
	}

	return nil
}

// CreateTimestampFile 在设备根目录创建一个以复制时间命名的txt文件
func CreateTimestampFile(deviceID string) error {
	// 获取当前时间并格式化
	timestamp := time.Now().Format("2006-01-02_15-04-05")
	fileName := filepath.Join(deviceID, timestamp+".txt")

	// 写入文件内容
	content := fmt.Sprintf("文件夹复制时间: %s\n", timestamp)
	if err := os.WriteFile(fileName, []byte(content), 0644); err != nil {
		return fmt.Errorf("创建时间戳文件失败: %v", err)
	}

	return nil
}

// ProcessDevice 处理单个存储设备，包括可选的删除所有文件和复制文件夹
func ProcessDevice(sourceFolder, deviceID string, deleteFirst bool, successCount, failCount *int, mutex *sync.Mutex, wg *sync.WaitGroup) {
	defer wg.Done()

	fmt.Printf("\n正在处理设备: %s\n", deviceID)

	// 如果需要先删除设备上的所有文件
	if deleteFirst {
		fmt.Printf("正在删除设备 '%s' 上的所有文件...\n", deviceID)
		if err := DeleteAllFilesOnDevice(deviceID); err != nil {
			fmt.Printf("删除设备文件失败: %v\n", err)
			mutex.Lock()
			*failCount++
			mutex.Unlock()
			return
		}
	}

	// 复制文件夹到设备
	fmt.Printf("正在复制文件夹 '%s' 到设备 '%s'...\n", sourceFolder, deviceID)
	if err := CopyFolderToDevice(sourceFolder, deviceID); err != nil {
		fmt.Printf("复制文件夹失败: %v\n", err)
		mutex.Lock()
		*failCount++
		mutex.Unlock()
		return
	}

	// 创建时间戳文件
	fmt.Printf("正在创建时间戳文件...\n")
	if err := CreateTimestampFile(deviceID); err != nil {
		fmt.Printf("创建时间戳文件失败: %v\n", err)
		mutex.Lock()
		*failCount++
		mutex.Unlock()
		return
	}

	fmt.Printf("设备 %s 处理完成!\n", deviceID)
	mutex.Lock()
	*successCount++
	mutex.Unlock()
}

func main() {
	// 解析命令行参数
	deleteFlag := flag.String("flag", "", "操作模式: delete表示先删除设备所有文件再复制")
	flag.Parse()

	// 确定是否需要先删除设备上的所有文件
	deleteFirst := *deleteFlag == "delete"

	// 源文件夹路径
	sourceFolder := "D:\\system\\document\\SD_PACKET\\private"

	// 检查源文件夹是否存在
	if _, err := os.Stat(sourceFolder); os.IsNotExist(err) {
		fmt.Printf("错误: 源文件夹 '%s' 不存在\n", sourceFolder)
		return
	}

	fmt.Println("正在获取连接的USB存储设备...")

	// 获取所有USB存储设备的设备ID
	devices, err := GetUSBDeviceIDs()
	if err != nil {
		fmt.Printf("获取USB设备失败: %v\n", err)
		return
	}

	// 初始化计数器
	totalDevices := len(devices)
	successCount := 0
	failCount := 0

	// 创建互斥锁和等待组，用于并发控制
	var mutex sync.Mutex
	var wg sync.WaitGroup

	if totalDevices == 0 {
		fmt.Println("未找到连接的USB存储设备")
	} else {
		if deleteFirst {
			fmt.Printf("找到 %d 个USB存储设备，开始多线程删除所有文件后再复制...\n", totalDevices)
		} else {
			fmt.Printf("找到 %d 个USB存储设备，开始多线程同时复制...\n", totalDevices)
		}

		// 为每个设备启动一个goroutine来并发处理
		wg.Add(totalDevices)
		for _, deviceID := range devices {
			go ProcessDevice(sourceFolder, deviceID, deleteFirst, &successCount, &failCount, &mutex, &wg)
		}

		// 等待所有goroutine完成
		fmt.Println("所有任务已启动，等待完成...")
		wg.Wait()
	}

	// 输出统计信息
	fmt.Printf("\n操作统计:\n")
	fmt.Printf("共处理 %d 个存储设备\n", totalDevices)
	fmt.Printf("成功: %d 个\n", successCount)
	fmt.Printf("失败: %d 个\n", failCount)
}
